# from src.dataset.data import  get_training_set
import os
from os.path import join


import numpy as np
from PIL import Image
import torch.nn as nn
from math import log10

def psnr(sr, hr):
    criterion = nn.MSELoss()
    mse = criterion(sr, hr)
    psnr = 10 * log10(1 / mse)
    return psnr

if __name__ == '__main__':
    # context.set_context(mode=context.PYNATIVE_MODE, device_id=6,
    #                     device_target="Ascend")
    SR_path = r'G:\Project\Python\AI\DBPN-Pytorch-master\Results\L4SR'
    HR_path = r'C:\Users\yan\Desktop\DBPN\image\Set5_LR_x4\SR'
    HR_imgs = sorted(os.listdir(HR_path))
    SR_imgs = sorted(os.listdir(SR_path))
    sr_img_files = [join(SR_path, x) for x in SR_imgs]
    hr_img_files = [join(HR_path, x) for x in HR_imgs]
    assert len(sr_img_files) == len(hr_img_files)
    psnr_sum = 0
    for i in range(len(hr_img_files)):
        sr_img = Image.open(sr_img_files[i]).convert("RGB")
        hr_img = Image.open(hr_img_files[i]).convert('RGB')
        sr_img = np.array(sr_img) / 255
        hr_img = np.array(hr_img) / 255
        # sr_img = Tensor(sr_img, dtype=mstype.float32)
        # hr_img = Tensor(hr_img, dtype=mstype.float32)
        tmp_psnr = psnr(sr_img, hr_img)
        print('the {%d/%d} image tmp_psnr = %.4f'%(i+1, len(hr_img_files),
                                                   tmp_psnr))
        psnr_sum += tmp_psnr
    print(psnr_sum / len(hr_img_files))